import PyPDF2

# 用户输入PDF文件
print('请输入pdf文档完整路径（例子：E:\BaiduNetdiskDownload\你好.pdf）：', end='')
pdf_path = input()

#以二进制形式打开文档
open_pdf = open(pdf_path, 'rb')
#读取文档的数据
read_pdf = PyPDF2.PdfReader(open_pdf)
#获取文档的总页数
get_pdf_pages = len(read_pdf.pages)

# 创建一个新的PDF文件以保存加密后的文档
pdf_out = PyPDF2.PdfWriter()

# 遍历源PDF的每一页，并将它们添加到新的PDF中
for get_pdf_page in range(get_pdf_pages):
    page = read_pdf.pages[get_pdf_page]
    pdf_out.add_page(page)

#提示用户输入密码
print('请输入密码：', end='')
password = str(input())

# 应用加密和权限设置
pdf_out.encrypt(password)

# 获取用户输入的文件名，不包括路径
pdf_file_name = pdf_path.split('\\')[-1]  # 使用 \\
pdf_base_name = pdf_file_name.split('.')[0]  # 去掉扩展名
# 指定保存 .pdf 文件的完整路径和名称
output_dir = "D:/"  # 替换为您的实际输出目录
output_file_name = f"{pdf_base_name}_加密.pdf"
output_path = output_dir + output_file_name
pdf_out.write(output_path)
pdf_out.close()
print('加密好了，快去看一下吧。')
